package com.lisp.algorithm.select;

import java.util.Arrays;

import com.lisp.algorithm.util.SortUtil;

public class MedianSelect {

    public static void main(String[] args) {
        SortUtil.test(MedianSelect::test);
    }
    
    public static void test(int a[]) {
        int copy[] = Arrays.copyOf(a, a.length);
        for(int i=1; i<=copy.length; i++) { 
            int index = medianSelect(copy, i);
            a[i-1] = copy[index];
        }
    }
    
    public static int medianSelect(int a[], int ith) {
        return medianSelectTailOptimise(a, ith, 0, a.length);
    }
    
    private static int medianSelectTailOptimise(int a[], int ith, int start, int end) {
        while(start < end-1) {            
            int pivotIndex = pivotIndex(a, ith, start, end);
            int p = partition(a, start, end, pivotIndex, ith);
            int k = p-start+1;
            if(k == ith) {
                return p;
            } else if(k < ith) {
                ith = ith - k;
                start = p+1;
            } else {
                end = p;
            }
        }
        
        return start;
    }
    
    private static int medianSelect(int a[], int ith, int start, int end) {
        if(start < end-1) {
            int pivotIndex = pivotIndex(a, ith, start, end);
            int p = partition(a, start, end, pivotIndex, ith);
            int k = p-start+1;
            if(k == ith) {
                return p;
            } else if(k < ith) {
                return medianSelect(a, ith-k, p+1, end);
            } else {
                return medianSelect(a, ith, start, p);
            }
        } else {
            return start;
        }
    }

    private static int pivotIndex(int[] a, int ith, int start, int end) {
        if(start > end-5) {
            return partition5(a, start, end);
        }
        
        for(int i=start; i<end; i=i+5) {
            int subEnd = i + 5;
            if(subEnd > end) {
                subEnd = end;
            }
            
            int mid = partition5(a, i, subEnd);
            SortUtil.swap(a, mid, (mid-start)/5 + start);
        }
        
        return medianSelect(a, ith/5+1, start, start + (end-start) / 5);
    }

    private static int partition(int[] a, int start, int end, int pivotIndex, int ith) {
        int lessIndex = start - 1;
        int pivot = a[pivotIndex];
        SortUtil.swap(a, pivotIndex, end-1);
        for(int i=start; i<end-1; i++) {
            if(a[i] < pivot) {
                lessIndex++;
                SortUtil.swap(a, i, lessIndex);
            }
        }
        
        int equalIndex = lessIndex;
        for(int i=lessIndex+1; i<end-1; i++) {
            if(a[i] == pivot) {
                equalIndex++;
                SortUtil.swap(a, i, equalIndex);
            }
        }
        
        SortUtil.swap(a, equalIndex+1, end-1);
        
        int ithElementIndex = ith+start-1;
        if(ithElementIndex <= lessIndex) {
            return lessIndex+1;
        } else if(ithElementIndex <= equalIndex+1){
            return ith+start-1;
        } else {
            return equalIndex+1;
        }
    }
    
    private static int partition5(int a[], int start, int end) {
        for(int i=start+1; i<end; i++) {
            int insertVal = a[i];
            int j = i;
            while(j > start && a[j-1] > insertVal) {
                a[j] = a[j-1];
                j--;
            }
            a[j] = insertVal;
        }
        
        return (start + end) / 2;
    }
}
